package com.gitee.xiao.check;

/**
 * @author wang xiao
 * date 2024/7/10
 */
public class FaceDetectionHandler extends TextWebSocketHandler  {

    private static final List<WebSocketSession> sessions = new ArrayList<>();

    static {
        System.loadLibrary(Core.NATIVE_LIBRARY_NAME);
    }

    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        sessions.add(session); // 连接建立后添加会话
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception {
        sessions.remove(session); // 连接关闭时移除会话
    }

    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {
        // 处理收到的消息并广播给所有连接的会话
        for (WebSocketSession webSocketSession : sessions) {
            webSocketSession.sendMessage(message);
        }
    }

    // 推送人脸检测结果
    public void sendFaceDetectionResult(String imageBase64) {
        for (WebSocketSession session : sessions) {
            try {
                session.sendMessage(new TextMessage(imageBase64)); // 发送消息
            } catch (IOException e) {
                e.printStackTrace();
            }
        }
    }

    // 处理视频流，检测人脸
    public void processVideoStream() {
        VideoCapture camera = new VideoCapture(0);
        if (!camera.isOpened()) {
            System.out.println("Camera Error");
            return;
        }
        CascadeClassifier faceDetector = new CascadeClassifier("haarcascade_frontalface_alt.xml");
        Mat frame = new Mat();
        while (camera.read(frame)) {
            Mat frameGray = new Mat();
            Imgproc.cvtColor(frame, frameGray, Imgproc.COLOR_BGR2GRAY);
            Imgproc.equalizeHist(frameGray, frameGray);

            Rect[] facesArray = faceDetector.detectMultiScale(frameGray);
            for (Rect face : facesArray) {
                Imgproc.rectangle(frame, face.tl(), face.br(), new Scalar(0, 255, 0), 3);
            }

            BufferedImage image = matToBufferedImage(frame);
            String imageBase64 = imageToBase64(image);
            sendFaceDetectionResult(imageBase64);
        }
        camera.release();
    }

    private BufferedImage matToBufferedImage(Mat mat) {
        // Convert Mat to BufferedImage
        MatOfByte mob = new MatOfByte();
        Imgcodecs.imencode(".jpg", mat, mob);
        byte[] byteArray = mob.toArray();
        BufferedImage bufImage = null;
        try {
            bufImage = ImageIO.read(new ByteArrayInputStream(byteArray));
        } catch (IOException e) {
            e.printStackTrace();
        }
        return bufImage;
    }

    private String imageToBase64(BufferedImage image) {
        // Convert BufferedImage to Base64 String
        ByteArrayOutputStream baos = new ByteArrayOutputStream();
        try {
            ImageIO.write(image, "jpg", baos);
            byte[] bytes = baos.toByteArray();
            return Base64.getEncoder().encodeToString(bytes);
        } catch (IOException e) {
            e.printStackTrace();
        }
        return null;
    }
}
}
